import numpy as np
import open3d as o3d
import cv2
# 这个文件read_write_model的引用来自于colmap的源码中文件
from read_write_model import read_cameras_binary, read_images_binary
import os

def __get_rotation_matrix__(yaw_deg):
    '''
    获取旋转指定角度的旋转矩阵。
    
    比如当前物体，朝向是[0,0,1]，沿着朝向是Z轴正方向，如果想要让它朝X轴正反向看，就要让它绕着Y轴旋转90°
    而这个让从当前朝向，旋转到目标朝向的矩阵，就是旋转矩阵
    '''
    # np.radians，用于实现将角度转化为弧度
    yaw = np.radians(yaw_deg)
    # 返回旋转矩阵
    return np.array([
        [np.cos(yaw), 0, np.sin(yaw)],
        [0, 1, 0],
        [-np.sin(yaw), 0, np.cos(yaw)]
    ])

def simulate_drive_view(mesh_path, images_bin, cameras_bin, image_dir, 
                        current_position, view_yaw_deg, fov_deg=45):
    
    '''
    基于图片相似度，返回最接近角度和位置的图片
    '''
    # 加载 mesh
    mesh = o3d.io.read_triangle_mesh(mesh_path)
    # 计算顶点的法线(normal)
    mesh.compute_vertex_normals()
    # 获取mesh的顶点（此时的顶点已经包含了发现信息
    vertices = np.asarray(mesh.vertices)

    # 加载图像和相机
    images = read_images_binary(images_bin)
    cameras = read_cameras_binary(cameras_bin)

    # 当前方向向量（前进方向，考虑偏航角）
    forward = np.array([0, 0, 1])  # 假设默认前方为 Z+
    # 获取平移30°角（view_yaw_deg）的旋转矩阵
    rot = __get_rotation_matrix__(view_yaw_deg)
    # 计算获得旋转之后的位置信息
    direction = rot @ forward

    # 初始化计算变量
    best_score = -1
    best_image_id = None
    # 遍历从image.bin文件中加载的文件
    for image_id, img in images.items():
        cam = cameras[img.camera_id]
        # 完成世界坐标系到相机坐标系
        center = -img.qvec2rotmat().T @ img.tvec  # camera center
        # [2, :]是相机坐标系的Z轴方向（也就是镜头的正前方），即Z+方向
        # 但是在colmap中Z+方向是“往后看”，“*-1”代表“向前看”
        view_dir = img.qvec2rotmat()[2, :] * -1   # camera viewing direction
        
        # 计算你当前位置（入参，调用时指定）指向相机的向量
        vec_to_cam = center - current_position
        vec_to_cam /= np.linalg.norm(vec_to_cam)
        # 将当前图片和相机的角度以及目标角度，计算相似度（cos夹角）
        score = np.dot(vec_to_cam, direction)
        # 根据相似度，来更新最佳图片id
        if score > best_score:
            best_score = score
            best_image_id = image_id
    # 如果有符合的图片，则通过cv2.imread读取并展示出来
    if best_image_id is not None:
        img = images[best_image_id]
        img_path = os.path.join(image_dir, img.name)
        image = cv2.imread(img_path)
        return image
        print(f"Selected image: {img.name}, score: {best_score:.2f}")
    else:
        print("No suitable image found.")
        return None

workspace= "/data/workspace/3D/room_ws"
# 示例使用
image = simulate_drive_view(
    mesh_path = os.path.join(workspace, "dense/0/meshed-poisson.ply"),
    images_bin = os.path.join(workspace, "sparse/0/images.bin"),
    cameras_bin = os.path.join(workspace, "sparse/0/cameras.bin"),
    image_dir = os.path.join(workspace, "images"),
    current_position=np.array([0, 0, 0]),  # 起始点，可自由设定
    view_yaw_deg=100                        # 向右看30°
)

if image is not None:
    cv2.imshow("Simulated Drive View", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
